import pysam
import argparse
import sys
import os
from multiprocessing import Manager
import multiprocessing as mp
import gzip
import mmap
import pod5 as p5
import time
import datetime
import numpy as np
from ont_fast5_api.fast5_interface import get_fast5_file
from pathlib import Path
from deepsignal3.utils.process_utils import CIGAR2CODE
from deepsignal3.utils.process_utils import CIGAR_REGEX
from deepsignal3.utils.process_utils import complement_seq
from deepsignal3.extract_features_pod5 import _group_signals_by_movetable_v2
import re
#from memory_profiler import profile

def get_object_size(obj):
    return sys.getsizeof(obj)
def parse_cigar(r_cigar, strand, ref_len):
    fill_invalid = -1
    # get each base calls genomic position
    r_to_q_poss = np.full(ref_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar if strand == 1 else r_cigar[::-1]
    for op_len, op in cigar_ops:
        if op == 1:
            # inserted bases into ref
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            for r_pos in range(curr_r_pos, curr_r_pos + op_len):
                r_to_q_poss[r_pos] = curr_q_pos
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                r_to_q_poss[curr_r_pos + op_offset] = curr_q_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    r_to_q_poss[curr_r_pos] = curr_q_pos
    if r_to_q_poss[-1] == fill_invalid:
        raise ValueError((
            'Invalid cigar string encountered. Reference length: {}  Cigar ' +
            'implied reference length: {}').format(ref_len, curr_r_pos))

    return r_to_q_poss
def get_q2tloc_from_cigar(r_cigar_tuple, strand, seq_len):
    """
    insertion: -1, deletion: -2, mismatch: -3
    :param r_cigar_tuple: pysam.alignmentSegment.cigartuples
    :param strand: 1/-1 for fwd/rev
    :param seq_len: read alignment length
    :return: query pos to ref pos
    """
    fill_invalid = -2
    # get each base calls genomic position
    q_to_r_poss = np.full(seq_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar_tuple if strand == 1 else r_cigar_tuple[::-1]
    for op, op_len in cigar_ops:
        if op == 1:
            # inserted bases into ref
            for q_pos in range(curr_q_pos, curr_q_pos + op_len):
                q_to_r_poss[q_pos] = -1
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                q_to_r_poss[curr_q_pos + op_offset] = curr_r_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    q_to_r_poss[curr_q_pos] = curr_r_pos
    if q_to_r_poss[-1] == fill_invalid:
        raise ValueError(('Invalid cigar string encountered. Reference length: {}  Cigar ' +
                          'implied reference length: {}').format(seq_len, curr_r_pos))
    return q_to_r_poss

def get_refloc_of_methysite_in_motif(seqstr, motifset, methyloc_in_motif=0):
    """

    :param seqstr:
    :param motifset:
    :param methyloc_in_motif: 0-based
    :return:
    """
    motifset = set(motifset)
    strlen = len(seqstr)
    motiflen = len(list(motifset)[0])
    sites = []
    for i in range(0, strlen - motiflen + 1):
        if seqstr[i:i + motiflen] in motifset:
            sites.append(i+methyloc_in_motif)
    return sites
def find_all_c_positions(text):
    positions = []
    pos = text.find('C')
    while pos != -1:
        positions.append(pos)
        pos = text.find('C', pos + 1)
    return positions

def get_key2lines(feafile):
    if feafile.endswith(".gz"):
        infile = gzip.open(feafile, 'rt')
    else:
        infile = open(feafile, 'r')
    key2lines = {}
    lcnt = 0
    for line in infile:
        words = line.strip().split("\t")
        key=(words[4],words[0],words[2])#,words[0],words[2],words[1]
        if key not in key2lines:
            key2lines[key] = []
        key2lines[key].append(lcnt) 
        lcnt += 1
    print('build index for tsv')
    infile.close()
    return key2lines

def _convert_cigarstring2tuple(cigarstr):
    # tuple in (oplen, op) format, like 30M -> (30, 0)
    return [(int(m[0]), CIGAR2CODE[m[-1]]) for m in CIGAR_REGEX.findall(cigarstr)]

def read_signal_bam(output_Q,filenames,bam,key2lines,features_file,is_pod5=False,motif='C',mod_loc=0, d_batch_size=2,output_size=20):
    
    fea_list=[]
    bam_file=pysam.AlignmentFile(bam,'rb',check_sq=False)
    print('%s: Building BAM index.' %str(datetime.datetime.now()), flush=True)
    bam_index=pysam.IndexedReads(bam_file)
    bam_index.build()
    print('%s: Finished building BAM index.' %str(datetime.datetime.now()), flush=True)
    file_type= 'pod5' if is_pod5 else 'fast5'
    signal_files= [filenames] if os.path.isfile(filenames) else Path(filenames).rglob("*.%s" %file_type)
    keep_num=0
    skip_num=0
    read_num=0
    signal_num=0
    seq_num=0
    signal_file_num=0
    if is_pod5:
        for filename in signal_files:
            signal_file_num+=1
            with p5.Reader(filename) as reader:
                print('Begin read pod5.')
                for read in reader.reads():
                    if output_Q.qsize()>=output_size:
                        time.sleep(20)
                        print('Pausing input due to INPUT queue size limit. Signal_qsize=%d' %(output_Q.qsize()), flush=True)
                    read_name=str(read.read_id)
                    try:
                        signal_num+=1
                        read_iter=bam_index.find(read_name)
                        seq_num+=1
                        for bam_read in read_iter:
                            read_num+=1
                            if bam_read.is_supplementary or bam_read.is_secondary or bam_read.is_mapped==False:
                                skip_num+=1
                                continue
                            read_dict=dict(bam_read.tags)
                            signal=read.signal
                            #strand = "-" if bam_read.is_reverse else "+"

                            seq = bam_read.query_sequence
                            if seq is None or signal is None:
                                skip_num+=1
                                continue
                            strand_code = 0 if bam_read.is_reverse else 1
                            reference_name=bam_read.reference_name
                            strand = "-" if strand_code==0 else "+"
                            
                            find_key=(read_name,reference_name,strand)
                            #if find_key not in key2lines:
                            #    print('can\'t find read {} in deepsignal'.format(read_name))
                            #    skip_num+=1
                            #    continue
                            qalign_start = bam_read.query_alignment_start
                            qalign_end = bam_read.query_alignment_end
                            ref_start=bam_read.reference_start
                            ref_end=bam_read.reference_end
                            #if bam_read.is_reverse:
                            #    seq_start = len(seq) - qalign_end
                            #    seq_end = len(seq) - qalign_start
                            #else:
                            #    seq_start = qalign_start
                            #    seq_end = qalign_end
                            
                            #cigartuples=bam_read.cigartuples
                            #seq_len=(seq_end - seq_start)
                            #try:
                            #readlocs = get_q2tloc_from_cigar(cigartuples, strand_code,seq_len)#q_to_r_poss
                            #    #num_bases = (seq_len - 1) // 2
                            #    #if bam_read.is_reverse:
                            #    #    seq = complement_seq(seq)
                            #    
                            #    offsets = find_all_c_positions(seq)
                            #    #rev_offset_loc = (len(motif) - 1 - mod_loc) - mod_loc
                            #    pred_deepsignal=[]
                            # 
                            #    find_first=0
                            #    read_loc=[]
                            #    
                            #
                            #    for off_loc_i in range(len(offsets)):
                            #        off_loc = offsets[off_loc_i]
                            #        #rev_loc = off_loc + rev_offset_loc
                            #        #rev_loc_in_rev = len(seq) - 1 - rev_loc
                            #        if seq_start <= off_loc < seq_end:
                            #            offset_idx = off_loc - seq_start
                            #            #offset_revidx = rev_loc - seq_start
                            #            if readlocs[offset_idx] != -1:
                            #                if bam_read.is_reverse:
                            #                    read_loc.append(ref_end - 1 - readlocs[offset_idx]) 
                            #                else:
                            #                    read_loc.append(readlocs[offset_idx] + ref_start)                          
                            #except Exception as e:
                            #    print('convert read loc false.', flush=True)
                            #    print(e, flush=True)
                            #consider C+h?;C+m?;
                            pattern = re.compile(r'C[+-]m')
                            matches = pattern.findall(read_dict["MM"])
                            if not matches:
                                print("skip methylation isn\'t 5mC")
                                continue
                            dorado_pred = []
                            dorado_pred_pos=[]
                            query_pos=[]
                            ref_pos=[]
                            if bam_read.modified_bases != None:
                                ref_loc = bam_read.get_reference_positions(full_length=True)
                                for m, locs in bam_read.modified_bases.items():
                                    if m[0] == 'C' and m[2] == 'm':
                                        for lc in locs:
                                            if ref_loc[lc[0]] != None:
                                                query_pos.append(lc[0] if bam_read.is_forward else (lc[0]-1))  
                                                rloc = ref_loc[lc[0]]  # if bam_read.is_forward else (ref_start+ref_end-int(ref_loc[lc[0]]))
                                                dorado_pred.append(lc[1])
                                                dorado_pred_pos.append(rloc)
                                                #if bam_read.is_reverse:
                                                #    ref_pos.append(ref_end - 1 - int(readlocs[lc[0]]))
                                                #else:
                                                #    ref_pos.append(ref_start + int(readlocs[lc[0]]))
                            flag=0
                            for p in query_pos:
                                if seq[p]!='C':
                                    print('-----------')
                                    print(strand)
                                    print(seq)
                                    print(','.join([str(x) for x in query_pos]),flush=True)
                                    flag=1
                                    break
                            if flag==1:
                                continue
                            #if bam_read.modified_bases != None:
                            #    for m, locs in bam_read.modified_bases_forward.items():
                            #        if m[0] == 'C' and m[2] == 'm':
                            #            for lc in locs:
                            #                query_pos.append(lc[0])
                            #if read_dict["MM"][0]!='C':
                            #    print("skip methylation isn\'t C")
                            #    continue
                            #if read_dict["MM"][2]!='m' and (str(read_dict["MM"][2])+str(read_dict["MM"][3]))!='mh':
                            #    print("skip methylation isn\'t m")
                            #    continue
                            #if len(read_dict["MM"].strip(';').split(','))<=1:
                            #    continue
                            #mm=read_dict["MM"].strip(';').split(',')[1:]
                            
                            #dorado_pred_pos=[]
                            #accum_shift=-1
                            #for shift in mm:
                            #    try:
                            #        accum_shift+=int(shift)+1
                            #        if accum_shift>=len(read_loc):
                            #            break
                            #    except Exception as e:
                            #        print(e)
                            #        #print(read_dict["MM"][0])
                            #        #print(read_dict["MM"][1])
                            #        #print(read_dict["MM"][6])
                            #        #print(read_dict["MM"][7])
                            #        #print(read_dict["MM"][8])
                            #    try:
                            #        dorado_pred_pos.append(read_loc[accum_shift])
                            #    except Exception as e:
                            #        print(len(read_loc))
                            #        print(accum_shift)
                            #        print(seq[accum_shift])
                            #        print(strand)
                            #        for c in offsets:
                            #            print(seq[c], end="")
                            #        print('\n')
                            #        print(e)
                            line_index=0
                            dorado_len=len(dorado_pred_pos)
                            #pred_num=0
                            #intersection_c=[]
                            pred_deepsignal={}
                            find_first=0  
                            pred_deepsignal_pos=[]
                            if features_file.endswith(".gz"):
                                infile = gzip.open(features_file, 'rt')
                            else:
                                infile = open(features_file, 'r')
                            for line in infile:
                                
                                # 在字节块的开始位置搜索换行符
                                #line = mm.readline()
                                #if line_index not in key2lines[find_key]:
                                #    if find_first==1:
                                #        break
                                #    else:
                                #        line_index+=1
                                #        continue
                                    #try:
                                    #    decoded_line = line.decode('cp1252')  # 尝试使用 utf-8 编码解码
                                    #except UnicodeDecodeError:
                                    #    decoded_line = line.decode('latin-1')
                                words=line.strip().split("\t")
                                ikey=(words[4],words[0],words[2])
                                if ikey!=find_key:
                                    continue
                                pred_deepsignal_pos.append(int(words[1]))
                                if int(words[1]) in dorado_pred_pos:
                                    pred_deepsignal[int(words[1])]=float(words[7])
                                #pred_num+=1
                                #line_index+=1
                                #if find_first==0:
                                #    find_first=1
                            #with mmap.mmap(infile.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                            #    for line in mm.readline():
                            #        # 在字节块的开始位置搜索换行符
                            #        #line = mm.readline()
                            #        if line_index not in key2lines[find_key]:
                            #            if find_first==1:
                            #                break
                            #            else:
                            #                line_index+=1
                            #                continue
                            #        try:
                            #            decoded_line = line.decode('cp1252')  # 尝试使用 utf-8 编码解码
                            #        except UnicodeDecodeError:
                            #            decoded_line = line.decode('latin-1')
                            #        words=decoded_line.strip().split("\t")
                            #        pred_deepsignal_pos.append(int(words[1]))
                            #        if int(words[1]) in dorado_pred_pos:
                            #            pred_deepsignal[words[1]]=float(words[7])
                            #        #pred_num+=1
                            #        line_index+=1
                            #        if find_first==0:
                            #            find_first=1
                            for key in dorado_pred_pos:
                                if key not in pred_deepsignal.keys():
                                    pred_deepsignal[key]=-1
                            #print('forward: {}'.format(bam_read.is_forward))
                            #print('reverse: {}'.format(bam_read.is_reverse))
                            #pred_deepsignal_text=','.join([str(pred_deepsignal[x]) for x in dorado_pred_pos])
                            #for i in range(dorado_len):
                            #    dorado_pred[i]=(dorado_pred[i]+0.5)/256
                            #pred_dorado_text=','.join([str(x) for x in dorado_pred])
                            #print("pred_dorado_text: ")
                            #print(pred_dorado_text)
                            #print("pred_deepsignal_pos: ")
                            #print(pred_deepsignal_pos)
                            #print(pred_deepsignal_text)
                            #print("dorado_pred_pos: ")
                            #print(dorado_pred_pos)
                            #print('ref_start: {}'.format(ref_start))
                            #print('ref_end: {}'.format(ref_end))
                            #print('qalign_start: {}'.format(qalign_start))
                            #print('qalign_end: {}'.format(qalign_end))
                            #print('seq len:{}'.format(len(seq)))
                            ##print("ref_pos: ")
                            ##print(ref_pos)
                            #print(query_pos)
                            #print(find_key)
                            #print("key2lines: ")
                            #print(key2lines[find_key])
                            #for i in query_pos:
                            #    print(seq[i],end='')
                            #print('\n')
                            #ref_seq=bam_read.get_reference_sequence()
                            #for i in dorado_pred_pos:
                            #    print(ref_seq[i],end='')
                            #print('\n')

                            #break
                            try:
                                num_trimmed = read_dict["ts"]
                                norm_shift = read_dict["sm"]
                                norm_scale = read_dict["sd"]
                                mv_table=read_dict["mv"]
                                
                                if num_trimmed >= 0:
                                    signal_trimmed = (signal[num_trimmed:] - norm_shift) / norm_scale
                                else:
                                   signal_trimmed = (signal[:num_trimmed] - norm_shift) / norm_scale
                                sshift, sscale = np.mean(signal_trimmed), float(np.std(signal_trimmed))
                                if sscale == 0.0:
                                    norm_signals = signal_trimmed
                                else:
                                    norm_signals = (signal_trimmed - sshift) / sscale
                                signal_group = _group_signals_by_movetable_v2(norm_signals, np.asarray(mv_table[1:]), int(mv_table[0]))
                                norm_signals_text=','.join([str(x) for x in signal_group])
                                #for i in pred_deepsignal.keys():
                                #    pred_deepsignal[i]=min(np.round(pred_deepsignal[i]*256),255)
                                
                                pred_deepsignal_text=','.join([str(pred_deepsignal[x]) for x in dorado_pred_pos])
                                for i in range(dorado_len):
                                    dorado_pred[i]=(dorado_pred[i]+0.5)/256

                                #pred_dorado_all=read_dict["ML"]#.strip(';').split(',')
                                #dorado_i=0
                                #for pos in dorado_pred_pos:
                                #    if pos in intersection_c:
                                #        pred_dorado.append(pred_dorado_all[dorado_i])
                                #    dorado_i+=1
                                pred_dorado_text=','.join([str(x) for x in dorado_pred])
                                                         
                                mean_pred=[]
                                k=0
                                for key in dorado_pred_pos:
                                    if pred_deepsignal[key]!=-1:
                                        mean_pred.append((dorado_pred[k]+pred_deepsignal[key])/2)#.astype(int)
                                    else:
                                        mean_pred.append(dorado_pred[k])
                                    k+=1
                                pred_label=[]
                                for i in range(dorado_len):
                                    if mean_pred[i]>=0.5:
                                        pred_label.append(1)
                                    else:
                                        pred_label.append(0)
                                mean_pred_text=','.join([str(x) for x in mean_pred])
                                pred_label_text=','.join([str(x) for x in pred_label])
                                pred_pos=','.join([str(x) for x in query_pos])
                                sample_id='\t'.join([read_name,reference_name,str(ref_start)])
                                fea_str='\t'.join([sample_id,seq,norm_signals_text,pred_pos,pred_dorado_text,pred_deepsignal_text,mean_pred_text,pred_label_text,str(bam_read.mapping_quality)])
                                keep_num+=1
                                fea_list.append(fea_str)
                            except Exception as e:
                                print('deal feature false.', flush=True)
                                #print(read_dict["ML"])
                                #print(read_dict.keys())
                                print(e, flush=True)
                        if len(fea_list)>=d_batch_size:
                           #print('output_Q++')
                           output_Q.put(fea_list)
                           fea_list=[]
                    except KeyError as e:
                        print('can\'t find read {} in bam.'.format(read_name), flush=True)
                        print(e, flush=True)
                        continue
                if len(fea_list)>0:
                    output_Q.put(fea_list)
    else:
        for filename in signal_files:
            with get_fast5_file(filename, mode="r") as f5:
                for read in f5.get_reads():
                    if output_Q.qsize()>2:
                        time.sleep(20)
                    read_name=read.read_id
                    try:
                        read_iter=bam_index.find(read_name)
                        for bam_read in read_iter:
                            read_dict=bam_read.to_dict()
                            signal=read.get_raw_data()
                            #strand = "-" if bam_read.is_reverse else "+"
                            qalign_start = bam_read.query_alignment_start
                            qalign_end = bam_read.query_alignment_end
                            seq = bam_read.query_sequence
                            reference_name=bam_read.reference_name
                            strand_code = -1 if bam_read.is_reverse else 1
                            strand = "-" if strand_code==-1 else "+"
                            #ref_start=bam_read.reference_start
                            #ref_end=bam_read.reference_end
                            if seq is None or signal is None:
                                skip_num+=1
                                continue
                            #item_cigar=bam_read.cigarstring
                            #cigartuple = _convert_cigarstring2tuple(item_cigar)
                            #if strand == "+" and cigartuple[-1][-1] == 2:
                            #    oplen_tmp = cigartuple[-1][0]
                            #    cigartuple = cigartuple[:-1]
                            #    ref_start -= oplen_tmp
                            #if strand == "-" and cigartuple[0][-1] == 2:
                            #    oplen_tmp = cigartuple[0][0]
                            #    cigartuple = cigartuple[1:]
                            #    ref_end += oplen_tmp
                            #if reference_name == 'read':
                            #    assert strand == "+"
                            #    r_to_q_poss = parse_cigar(cigartuple, strand_code, len(seq))
                            #else:
                            #    if strand == "-":
                            #        seq = complement_seq(seq)
                            #    r_to_q_poss = parse_cigar(cigartuple, strand_code, len(seq))
                            #ref_readlocs=[]
                            #for ref_pos, q_pos in enumerate(r_to_q_poss[:-1]):
                            #    ref_readlocs[ref_pos] = q_pos + qalign_start
                            ##cigartuples=bam_read.cigartuples
                            #seq_len=(qalign_end - qalign_start)
                            #seg_len_accum = 0
                            ##readlocs = get_q2tloc_from_cigar(cigartuples, strand_code,seq_len)
                            #num_bases = (seq_len - 1) // 2
                            #offsets = get_refloc_of_methysite_in_motif(seq, set(motif), mod_loc)
                            pred_deepsignal=[]
                            key=(read_name,reference_name,strand)
                            if key not in key2lines:
                                continue
                            find_first=0
                            line_index=0
                            with mmap.mmap(infile.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                                for line in mm.readline():
                                    # 在字节块的开始位置搜索换行符
                                    #line = mm.readline()
                                    if line_index!=key2lines[key]:
                                        if find_first==1:
                                            break
                                        line_index+=1
                                        continue
                                    try:
                                        decoded_line = line.decode('cp1252')  # 尝试使用 utf-8 编码解码
                                    except UnicodeDecodeError:
                                        decoded_line = line.decode('latin-1')
                                    words=decoded_line.strip().split()
                                    pred_deepsignal.append(str(words[6]))
                                    find_first=1
                                    line_index+=1
                            #for off_loc_i in range(len(offsets)):
                            #    off_loc = offsets[off_loc_i]
                            #    if num_bases <= off_loc < len(seq) - num_bases:
                            #        seg_len_accum += len(seq)                                  
                            #        if off_loc < seg_len_accum:
                            #            seg_off_loc = off_loc - (seg_len_accum - len(seq))
                            #        read_loc = ref_start + seg_off_loc if strand == "+" else ref_end - 1 - seg_off_loc
                            #        key=(read_name,reference_name,strand,read_loc)
                            #        if key not in key2lines:
                            #            continue
                            #        find_first=0
                            #        with mmap.mmap(infile.fileno(), 0, access=mmap.ACCESS_READ) as mm:
                            #            for line in mm.readline():
                            #            # 在字节块的开始位置搜索换行符
                            #            #line = mm.readline()
                            #                if line_index!=key2lines[key]:
                            #                    if find_first==1:
                            #                        break
                            #                    line_index+=1
                            #                    continue
                            #                try:
                            #                    decoded_line = line.decode('cp1252')  # 尝试使用 utf-8 编码解码
                            #                except UnicodeDecodeError:
                            #                    decoded_line = line.decode('latin-1')
                            #                words=decoded_line.strip().split()
                            #                pred_deepsignal.append(str(words[6]))
                            #                find_first==1
                            num_trimmed = read_dict["ts"]
                            norm_shift = read_dict["sm"]
                            norm_scale = read_dict["sd"]
                            mv_table=read_dict["mv"]#[2:]
                            if num_trimmed >= 0:
                                signal_trimmed = (signal[num_trimmed:] - norm_shift) / norm_scale
                            else:
                                signal_trimmed = (signal[:num_trimmed] - norm_shift) / norm_scale
                            sshift, sscale = np.mean(signal_trimmed), float(np.std(signal_trimmed))
                            if sscale == 0.0:
                                norm_signals = signal_trimmed
                            else:
                                norm_signals = (signal_trimmed - sshift) / sscale
                            norm_signals_text=','.join([str(x) for x in norm_signals])
                            pred_deepsignal_text=','.join([str(x) for x in pred_deepsignal])
                            pred_dorado=[]
                            for p in read_dict["ML"].split(','):
                                pred_dorado.append(float(p)/255)
                            pred_dorado_text=','.join([str(x) for x in pred_dorado])
                            fea_str='\t'.join([seq,norm_signals_text,mv_table,pred_dorado_text,pred_deepsignal_text])
                            fea_list.append(fea_str)
                            keep_num+=1
                        if len(fea_list)>=d_batch_size:
                           output_Q.put(fea_list)
                           fea_list=[]
                    except KeyError:
                        skip_num+=1
                        continue
                if len(fea_list)>0:
                    output_Q.put(fea_list)
    print('signal file number:{}; signal number: {}; sequence number: {}; read number: {};and skip number: {};and keep number: {}'.format(signal_file_num,signal_num,seq_num,read_num,skip_num,keep_num), flush=True)

def _write_featurestr(write_fp, featurestr_q,time_wait = 3):
    print('write_process-{} starts'.format(os.getpid()))
    output=0
    with open(write_fp, 'w') as wf:
        while True:
            # during test, it's ok without the sleep(time_wait)
            if featurestr_q.empty():
                time.sleep(time_wait)
                continue
            features_str = featurestr_q.get()
            if features_str == "kill":
                print('output line {}'.format(output))
                print('write_process-{} finished'.format(os.getpid()))
                break
            for one_features_str in features_str:
                output+=1
                wf.write(one_features_str + "\n")
            wf.flush()

def extract(args):
    manager = mp.Manager()
    output_Q = manager.Queue()
    input_path = os.path.abspath(args.tsv)
    if not os.path.exists(input_path):
        raise ValueError("--tsv input_path does not exist!")
    #total_size=os.path.getsize(input_path)
    #ex_size=int(total_size/ex_dp)
    key2lines=None#get_key2lines(input_path)
    print("size of key2lines: {}".format(get_object_size(key2lines)))
    sys.stdout.flush()
    pb=mp.Process(target=read_signal_bam,args=(output_Q,args.signal,args.bam,key2lines,input_path,args.pod5),name="signal_bam")
    pb.daemon = True
    pb.start()
    p_w = mp.Process(target=_write_featurestr, args=(args.write_path, output_Q),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    pb.join()
    output_Q.put("kill")
    p_w.join()



def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--bam", type=str,required=True)
    parser.add_argument("--pod5", action="store_true", default=False, required=False,
                       help='use pod5, default false')
    parser.add_argument("--signal", type=str,required=True)
    parser.add_argument("--tsv", type=str,required=True)
    parser.add_argument("--write_path", type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    parser.add_argument("--d_batch_size", action="store", type=int, default=2,
                         required=False)

    return parser.parse_args()


def main():
    args=parse_args()
    extract(args)


if __name__ == '__main__':
    sys.exit(main())